#include <kuai/sim/gromacs.hpp>

namespace kuai {
	
	bool GromacsTOPReader::update(std::istream& stream, PhaseSpace& space) const {

		std::map<String, boost::shared_ptr<PhaseSpace> > mols;
		String line;
		while (getline(stream, line)) {
			if (boost::algorithm::starts_with(line, "[ moleculetype ]")) {
				// std::pair<String, boost::shared_ptr<PhaseSpace> > typeI = parseMolType(stream);
			}
			else if (boost::algorithm::starts_with(line, "[ molecules ]")) {
			}
		}
		return true;
	}

	IMPLEMENT_WORKER_CREATOR(GromacsTrajecPlayer, gromacs_trr)

	GromacsTrajecPlayer::GromacsTrajecPlayer()
	{
		file = NULL;
	}

	void GromacsTrajecPlayer::open() { 
		file = fopen(filename.c_str(), "rb");
		if (file == NULL) {
			check_last_error();
			throw error("Can not open file %1%", filename);
		}
		nextTRN();
	}


	void GromacsTrajecPlayer::close() { 
		if (file == NULL) {
			throw error("File %1% is not opened yet", filename);
		}
		else {
			fclose(file);
			file = NULL;
			_trnHeader.step = -1;
		}
	}

	size_t GromacsTrajecPlayer::current() 
	{ 
		return _trnHeader.step;
	}

	void GromacsTrajecPlayer::update(PhaseSpace& space, Index tid, Index nThreads) {
		if (current() == size_t(-1)) {
			throw EndOfTrajec(filename);
		}
		
		if (space.countAtoms() == 0) {
			std::vector<Index> v(_trnHeader.natoms+2, 0);
			v[0] = _trnHeader.natoms;
			space[MOL_INFO_CT] = boost::shared_ptr<SimObjPtr>(new PropertyConnectTable(&v[0]));
		}
		else if (space.countAtoms() != _trnHeader.natoms) {
			throw error("The Traject File %1% is not fit for this molecule.", filename);
		}

		bool ok = true;
		int nskip = _trnHeader.ir_size+_trnHeader.e_size;
		if (nskip != 0) {
			fseek(file, nskip, SEEK_CUR);
		}
		if (_trnHeader.bDouble) {
			if (_trnHeader.box_size) {
				double box[9];
				ok = ok && readDouble(file, box, 9);
				XYZ a((RealNumber)box[0], (RealNumber)box[1], (RealNumber)box[2]);
				XYZ b((RealNumber)box[3], (RealNumber)box[4], (RealNumber)box[5]);
				XYZ c((RealNumber)box[6], (RealNumber)box[7], (RealNumber)box[8]);
				RealNumber alpha = angle(b, c);
				RealNumber beta = angle(a, c);
				RealNumber gamma = angle(a, b);

				PhaseSpace::const_iterator 
					it = space.find(MOL_INFO_PBC);
				PropertyPBC* p;
				if (it == space.end()) {
					p = new PropertyPBC;
					space[MOL_INFO_PBC] = boost::shared_ptr<SpaceProperty>(p);
				}
				else {
					p = static_cast<PropertyPBC*>((it->second).get());
				}
				*(p->x_ptr()) = PBC(a.abs()*10, b.abs()*10, c.abs()*10, alpha, beta, gamma);
			}
			nskip = _trnHeader.vir_size+_trnHeader.pres_size
				+_trnHeader.top_size+_trnHeader.sym_size;
			if (nskip != 0) {
				fseek(file, nskip, SEEK_CUR);
			}
			if (_trnHeader.x_size) {
				assert (_trnHeader.x_size == _trnHeader.natoms*3*sizeof(double));
				std::vector<double> v(_trnHeader.natoms*3);
				ok = ok && readDouble(file, &v[0], _trnHeader.natoms*3);
				PhaseSpace::const_iterator 
					it = space.find(MOL_INFO_ATOM_COORD_3D);
				PropertyXYZ* p;
				if (it == space.end()) {
					p = new PropertyXYZ(_trnHeader.natoms, XYZ(0, 0, 0));
					space[MOL_INFO_ATOM_COORD_3D] = boost::shared_ptr<SpaceProperty>(p);
				}
				else {
					p = static_cast<PropertyXYZ*>((it->second).get());
				}
				XYZ* coords=p->x_ptr();
				for (size_t i = 0; i < v.size(); i += 3) {
					coords[i/3].x = RealNumber(v[i]*10);
					coords[i/3].y = RealNumber(v[i+1]*10);
					coords[i/3].z = RealNumber(v[i+2]*10);
				}
			};
			if (_trnHeader.v_size) {
				assert (_trnHeader.x_size == _trnHeader.natoms*3*sizeof(double));
				std::vector<double> v(_trnHeader.natoms*3);
				ok = ok && readDouble(file, &v[0], _trnHeader.natoms*3);
				PhaseSpace::const_iterator 
					it = space.find(MOL_INFO_ATOM_SPEED);
				PropertyXYZ* p;
				if (it == space.end()) {
					p = new PropertyXYZ(_trnHeader.natoms, XYZ(0, 0, 0));
					space[MOL_INFO_ATOM_SPEED] = boost::shared_ptr<SpaceProperty>(p);
				}
				else {
					p = static_cast<PropertyXYZ*>((it->second).get());
				}
				XYZ* speeds=p->x_ptr();
				for (size_t i = 0; i < v.size(); i += 3) {
					speeds[i/3].x = RealNumber(v[i]/10);
					speeds[i/3].y = RealNumber(v[i+1]/10);
					speeds[i/3].z = RealNumber(v[i+2]/10);
				}
			}
			if (_trnHeader.f_size) {
				assert (_trnHeader.x_size == _trnHeader.natoms*3*sizeof(double));
				std::vector<double> v(_trnHeader.natoms*3);
				ok = ok && readDouble(file, &v[0], _trnHeader.natoms*3);
				PhaseSpace::const_iterator 
					it = space.find(MOL_INFO_ATOM_FORCE);
				PropertyXYZ* p;
				if (it == space.end()) {
					p = new PropertyXYZ(_trnHeader.natoms, XYZ(0, 0, 0));
					space[MOL_INFO_ATOM_FORCE] = boost::shared_ptr<SpaceProperty>(p);
				}
				else {
					p = static_cast<PropertyXYZ*>((it->second).get());
				}
				XYZ* forces=p->x_ptr();
				for (size_t i = 0; i < v.size(); i += 3) {
					forces[i/3].x = RealNumber(v[i]/100);
					forces[i/3].y = RealNumber(v[i+1]/100);
					forces[i/3].z = RealNumber(v[i+2]/100);
				}
			}
		}
		else {
			if (_trnHeader.box_size) {
				float box[9];
				ok = ok && readFloat(file, box, 9);
				XYZ a(box[0], box[1], box[2]);
				XYZ b(box[3], box[4], box[5]);
				XYZ c(box[6], box[7], box[8]);
				RealNumber alpha = angle(b, c);
				RealNumber beta = angle(a, c);
				RealNumber gamma = angle(a, b);

				PhaseSpace::const_iterator 
					it = space.find(MOL_INFO_PBC);
				PropertyPBC* p;
				if (it == space.end()) {
					p = new PropertyPBC;
					space[MOL_INFO_PBC] = boost::shared_ptr<SpaceProperty>(p);
				}
				else {
					p = static_cast<PropertyPBC*>((it->second).get());
				}
				*(p->x_ptr()) = PBC(a.abs()*10, b.abs()*10, c.abs()*10, alpha, beta, gamma);
			}
			nskip = _trnHeader.vir_size+_trnHeader.pres_size
				+_trnHeader.top_size+_trnHeader.sym_size;
			if (nskip != 0) {
				fseek(file, nskip, SEEK_CUR);
			}
			if (_trnHeader.x_size) {
				assert (_trnHeader.x_size == _trnHeader.natoms*3*sizeof(float));
				std::vector<float> v(_trnHeader.natoms*3);
				ok = ok && readFloat(file, &v[0], _trnHeader.natoms*3);
				PhaseSpace::const_iterator 
					it = space.find(MOL_INFO_ATOM_COORD_3D);
				PropertyXYZ* p;
				if (it == space.end()) {
					p = new PropertyXYZ(_trnHeader.natoms, XYZ(0, 0, 0));
					space[MOL_INFO_ATOM_COORD_3D] = boost::shared_ptr<SpaceProperty>(p);
				}
				else {
					p = static_cast<PropertyXYZ*>((it->second).get());
				}
				XYZ* coords=p->x_ptr();
				for (size_t i = 0; i < v.size(); i += 3) {
					coords[i/3].x = v[i]*10;
					coords[i/3].y = v[i+1]*10;
					coords[i/3].z = v[i+2]*10;
				}
			};
			if (_trnHeader.v_size) {
				assert (_trnHeader.x_size == _trnHeader.natoms*3*sizeof(float));
				std::vector<float> v(_trnHeader.natoms*3);
				ok = ok && readFloat(file, &v[0], _trnHeader.natoms*3);
				PhaseSpace::const_iterator 
					it = space.find(MOL_INFO_ATOM_SPEED);
				PropertyXYZ* p;
				if (it == space.end()) {
					p = new PropertyXYZ(_trnHeader.natoms, XYZ(0, 0, 0));
					space[MOL_INFO_ATOM_SPEED] = boost::shared_ptr<SpaceProperty>(p);
				}
				else {
					p = static_cast<PropertyXYZ*>((it->second).get());
				}
				XYZ* speeds=p->x_ptr();
				for (size_t i = 0; i < v.size(); i += 3) {
					speeds[i/3].x = v[i]/10;
					speeds[i/3].y = v[i+1]/10;
					speeds[i/3].z = v[i+2]/10;
				}
			}
			if (_trnHeader.f_size) {
				assert (_trnHeader.x_size == _trnHeader.natoms*3*sizeof(float));
				std::vector<float> v(_trnHeader.natoms*3);
				ok = ok && readFloat(file, &v[0], _trnHeader.natoms*3);
				PhaseSpace::const_iterator 
					it = space.find(MOL_INFO_ATOM_FORCE);
				PropertyXYZ* p;
				if (it == space.end()) {
					p = new PropertyXYZ(_trnHeader.natoms, XYZ(0, 0, 0));
					space[MOL_INFO_ATOM_FORCE] = boost::shared_ptr<SpaceProperty>(p);
				}
				else {
					p = static_cast<PropertyXYZ*>((it->second).get());
				}
				XYZ* forces=p->x_ptr();
				for (size_t i = 0; i < v.size(); i += 3) {
					forces[i/3].x = v[i]/100;
					forces[i/3].y = v[i+1]/100;
					forces[i/3].z = v[i+2]/100;
				}
			}
		}

		if (!ok) {
			throw error("Can not trajec file %1%", filename);
		}

		nextTRN();
	}

	bool GromacsTrajecPlayer::nextTRN() {
		static const int GROMACS_MAGIC = 0x000007C9;
		_reversed = false;
		int temp = 0;
		bool ok = readInteger(file, temp);	// Read the Magic Number
		if (temp == GROMACS_MAGIC) {
		}
		else if (temp == 0xC9070000) {
			_reversed = true;
		}
		else {
			ok = false;
		}
		ok = ok && readInteger(file, temp);	// Not know the meening.

		String version;
		ok = ok && readString(file, version);	// Read the version
		ok = ok && (version == "GMX_trn_file");

		RealNumber lasttime = _trnHeader.t;
		int laststep = _trnHeader.step;

		ok = ok && readInteger(file, _trnHeader.ir_size);	// discard
		ok = ok && readInteger(file, _trnHeader.e_size);	// discard
		ok = ok && readInteger(file, _trnHeader.box_size);	// box in 3*3 
		ok = ok && readInteger(file, _trnHeader.vir_size);	// discard
		ok = ok && readInteger(file, _trnHeader.pres_size); // discard
		ok = ok && readInteger(file, _trnHeader.top_size);	// discard
		ok = ok && readInteger(file, _trnHeader.sym_size);	// discard
		ok = ok && readInteger(file, _trnHeader.x_size);	// coordinates
		ok = ok && readInteger(file, _trnHeader.v_size);	// velocities
		ok = ok && readInteger(file, _trnHeader.f_size);	// forces

		ok = ok && readInteger(file, _trnHeader.natoms);	// number of atoms

		ok = (_trnHeader.x_size > 0);
		if (ok) {
			int realsize = _trnHeader.x_size / _trnHeader.natoms / 3;
			if (realsize == sizeof(double)) {
				_trnHeader.bDouble = true;
			}
			else if (realsize == sizeof(float)) {
				_trnHeader.bDouble = false;
			}
			else {
				throw error("Unknown Real Number Format");
			}

			ok = ok && readInteger(file, _trnHeader.step);
			ok = ok && readInteger(file, _trnHeader.nre);

			if (_trnHeader.bDouble) {
				double time, lambda;
				ok = ok && readDouble(file, time);
				ok = ok && readDouble(file, lambda);
				if (ok) {
					_trnHeader.t = RealNumber(time);
					_trnHeader.lambda = RealNumber(lambda);
				}
			}
			else {
				float time, lambda;
				ok = ok && readFloat(file, time);
				ok = ok && readFloat(file, lambda);
				if (ok) {
					_trnHeader.t = time;
					_trnHeader.lambda = lambda;
				}
			}
		}

		if (ok) {
			// stepLength = (_trnHeader.t - lasttime) / (_trnHeader.step - laststep) * 1000;
		}
		else {
			_trnHeader.step = -1;
		}
		return ok && !feof(file);
	}

	WORKER_HELP_FUNCTION(GromacsTrajecPlayer, output) {
		// TODO:
	}

}